"""
# Reference

[Deep Residual Learning for Image Recognition]()

This model is based on the following implementations:
- https://github.com/calmisential/TensorFlow2.0_ResNet
- https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py

"""
from tensorflow.keras import Model, layers, activations, Sequential


class BasicBlock(layers.Layer):
    def __init__(self, filter_num, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = layers.Conv2D(filter_num, 3, stride, 'same')
        self.bn1 = layers.BatchNormalization()
        self.relu = layers.ReLU()
        self.conv2 = layers.Conv2D(filter_num, 3, 1, 'same')
        self.bn2 = layers.BatchNormalization()
        if stride != 1:
            self.downsample = Sequential()
            self.downsample.add(layers.Conv2D(filter_num, 1, stride))
            self.downsample.add(layers.BatchNormalization())
        else:
            self.downsample = lambda x: x

    def call(self, inputs, training=None, **kwargs):
        identity = self.downsample(inputs)

        out = self.conv1(inputs)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)

        out += identity
        out = self.relu(out)
        return out


class Bottleneck(layers.Layer):
    def __init__(self, filter_num, stride=1):
        super(Bottleneck, self).__init__()
        self.conv1 = layers.Conv2D(filter_num, 1, 1, 'same')
        self.bn1 = layers.BatchNormalization()
        self.conv2 = layers.Conv2D(filter_num, 3, stride, 'same')
        self.bn2 = layers.BatchNormalization()
        self.conv3 = layers.Conv2D(filter_num * 4, 1, 1, 'same')
        self.bn3 = layers.BatchNormalization()
        self.relu = layers.ReLU()
        self.downsample = Sequential([
            layers.Conv2D(filter_num * 4, (1, 1), stride),
            layers.BatchNormalization()
        ])

    def call(self, inputs, training=None, **kwargs):
        identity = self.downsample(inputs)

        out = self.conv1(inputs)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        out += identity
        out = self.relu(out)
        return out


class ResNet(Model):
    def __init__(self, block, layer_params, num_classes=200):
        super(ResNet, self).__init__()
        self.conv1 = layers.Conv2D(64, 7, 2, 'same')
        self.bn1 = layers.BatchNormalization()
        self.relu = layers.ReLU()
        self.maxpool = layers.MaxPool2D(3, 2, 'same')
        self.layer1 = self._make_layer(block, 64, layer_params[0], 1)
        self.layer2 = self._make_layer(block, 128, layer_params[1], 2)
        self.layer3 = self._make_layer(block, 256, layer_params[2], 2)
        self.layer4 = self._make_layer(block, 512, layer_params[3], 2)
        self.avgpool = layers.GlobalAveragePooling2D()
        self.fc = layers.Dense(units=num_classes, activation=activations.softmax)

    @staticmethod
    def _make_layer(block, filter_num, block_num, stride=1):
        layer_seq = [block(filter_num, stride=stride)]
        for _ in range(1, block_num):
            layer_seq.append(block(filter_num, stride=1))
        return Sequential(layer_seq)

    def call(self, inputs, training=None, mask=None):
        x = self.conv1(inputs)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = self.fc(x)
        return x


def _resnet(block, layer_params):
    return ResNet(block, layer_params)


def resnet18():
    return _resnet(BasicBlock, [2, 2, 2, 2])


def resnet34():
    return _resnet(BasicBlock, [3, 4, 6, 3])


def resnet50():
    return _resnet(Bottleneck, [3, 4, 6, 3])
